#include "Util.h"

FnNtOpenProcess lpFnNtOpenProcess;

NTSTATUS NTAPI HkNtOpenProcess(
    _Out_ PHANDLE ProcessHandle,
    _In_ ACCESS_MASK DesiredAccess,
    _In_ POBJECT_ATTRIBUTES ObjectAttributes,
    _In_opt_ PCLIENT_ID ClientId
)
{
    NTSTATUS ntStatus;
    PEPROCESS lpEprocess;
    PCHAR lpName = NULL;

    ntStatus = PsLookupProcessByProcessId(ClientId->UniqueProcess, &lpEprocess);
    if (NT_SUCCESS(ntStatus))
    {
        lpName = PsGetProcessImageFileName(lpEprocess);
        if (lpName)
            DBGPRINT("ProcessName: %s\r\n", lpName);
    }

    return NtOpenProcess(ProcessHandle, DesiredAccess, ObjectAttributes, ClientId);
}

void __fastcall SysCallCallBack(unsigned int SystemCallIndex, void** SystemCallFunction)
{
    UNREFERENCED_PARAMETER(SystemCallIndex);

    if (*SystemCallFunction == lpFnNtOpenProcess)
    {
        *SystemCallFunction = (void*)HkNtOpenProcess;
    }
}

void HkSystemRoutine()
{
    UNICODE_STRING usNtOpenProcess = RTL_CONSTANT_STRING(L"NtOpenProcess");
    ULONG_PTR ulNtOpenProcess = (ULONG_PTR)MmGetSystemRoutineAddress(&usNtOpenProcess);

    if (ulNtOpenProcess)
        lpFnNtOpenProcess = (FnNtOpenProcess)ulNtOpenProcess;

    IfhInitialize(SysCallCallBack);
}

void HkCancel()
{
    IfhRelease();
}
